#!/usr/bin/python

import gdb
import struct
import numpy as np

class Mat:
    """ class Mat for cv::Mat data from GDB """

    __mat_types = {
        # depth : ( type, bytes, format)
        0 : ( 'CV_8U',  1,  'B'),
        1 : ( 'CV_8S',  1,  'b'),
        2 : ('CV_16U',  2,  'H'),
        3 : ('CV_16S',  2,  'h'),
        4 : ('CV_32S',  4,  'i'),
        5 : ('CV_32F',  4,  'f'),
        6 : ('CV_64F',  8,  'd')
    }

    def __init__(self, mat, inferior):
        assert mat.type.name == 'cv::Mat', 'Error input type( should be cv::Mat )!'
        assert inferior.is_valid(), 'Invalid gdb inferior!'
        self._flags = int(mat['flags'])
        self._dims = int(mat['dims'])
        self._rows = int(mat['rows'])
        self._cols = int(mat['cols'])
        self._data = int(mat['data']) # I don't know why it just return the address, however it works
        self._step = int(mat['step']['p'][0])
        self._step1 = int(mat['step']['p'][1])


        # get info from flags
        self._depth = self._flags & 7
        self._channels = 1 + (self._flags >> 3) & 0x1ff
        self._type = self.__mat_types[self._depth][0]

        # get data from memory
        byte_size = self._step * self._rows
        data_buffer = inferior.read_memory(self._data, byte_size)

        # convert bytes to its format
        bytes = self.__mat_types[self._depth][1]
        format = self.__mat_types[self._depth][2]
        fmt = '%i%s' % (byte_size / bytes, format)
        data_buffer = struct.unpack(fmt, data_buffer)
        if self._channels == 1:
            self._image = np.reshape(data_buffer, (self._rows, self._cols))
        elif self._channels == 3:
            self._image = np.reshape(data_buffer, (self._rows, self._cols, self._channels))

    def image(self):
        return self._image

    def __str__(self):
        return 'Mat object [dims:%d, type:%s, depth:%d, channels:%d, size:%dx%d, data:%s, step:(%d %d)]' % \
               (self._dims, self._type, self._depth, self._channels, self._rows, self._cols, \
                hex(self._data), self._step, self._step1)

class MatDetector:

    def __init__(self):
        self.mat_type = gdb.lookup_type('cv::Mat')

    def detect(self):
        frame = self._getGDBFrame()
        blocks = self._getGDBBlocks(frame)
        values = self._getGDBValues(blocks)
        inferior = gdb.selected_inferior()
        mats = self._getCVMats(values, inferior)
        return mats


    def _getGDBFrame(self):
        return gdb.newest_frame()

    def _getGDBBlocks(self, frame_):
        it = frame_.block().__iter__()
        blocks = []
        while True:
            try:
                x = next(it)
                #print(x)
                blocks.append(x)
            except StopIteration as e:
                break
        return blocks

    def _getGDBValues(self, blocks_):
        values = []
        for v in blocks_:
            # print(v)
            values.append(gdb.parse_and_eval(str(v)))
        return values

    def _getCVMats(self, values_, inferior_):
        mats = []
        for v in values_:
            if v.type.code == gdb.TYPE_CODE_REF:
                v = v.referenced_value()
            if v.type == self.mat_type:
                # print(v)
                mats.append([v, inferior_])
        return mats

